import asyncio
import os
import time
from typing import Dict, List, Any, Optional, Tuple, TypeVar, Generic, Union, Callable
from datetime import datetime
from abc import ABC, abstractmethod
import httpx
from app.utils.logger import get_logger

T = TypeVar('T')  # 查询参数类型
R = TypeVar('R')  # 结果数据类型

class AsyncDataService(Generic[T, R], ABC):
    """
    异步数据服务的抽象基类
    实现了通用的批处理和缓存逻辑，子类只需实现特定的数据获取和缓存操作
    
    泛型参数:
        T: 查询参数类型 (如 str 或 int)
        R: 结果数据类型 (通常是 Dict[str, Any])
    """
    
    def __init__(self, lang: str = "en", batch_size: int = 50, max_wait_time: float = 0.5, ):
        """
        初始化异步数据服务
        
        Args:
            batch_size: 批处理的最大任务数量
            max_wait_time: 最大等待时间（秒），即使未达到批处理大小也会处理任务
        """
        self.lang = lang
        self.task_queue = asyncio.Queue()  # 异步任务队列
        self.results = {}  # 存储任务结果的字典
        self.batch_size = batch_size
        self.max_wait_time = max_wait_time
        self.lock = asyncio.Lock()  # 用于异步安全访问结果字典
        self.event = asyncio.Event()  # 用于通知处理协程有新任务
        self.worker_task = None
        
        # 标记服务是否正在运行
        self.running = False
        
        # 获取日志记录器
        self.logger = get_logger(self.__class__.__name__)
        self.init_http_client()
        
    def init_http_client(self):
        self.httpx_client = httpx.AsyncClient(timeout=60.0, proxy=os.environ.get("HTTP_PROXY"))
    
    async def initialize(self):
        """
        异步初始化服务，启动后台任务
        """
        if not self.running:
            self.running = True
            # 启动后台处理协程
            self.worker_task = asyncio.create_task(self._process_tasks())
            self.logger.info("服务已初始化并启动")
        return self
    
    async def submit_task(self, query_param: T, cache_ttl: int, timeout: float = 360.0) -> R:
        """
        提交一个任务并等待结果
        
        Args:
            query_param: 查询参数 (如标题或ID)
            timeout: 等待结果的最大时间（秒）
            
        Returns:
            查询结果数据
            
        Raises:
            TimeoutError: 如果等待超时
        """
        # 首先检查缓存中是否已有数据
        self.logger.debug(f"提交任务: {query_param}")
        cached_data = await self._check_cache(query_param, cache_ttl)
        if cached_data:
            self.logger.debug(f"缓存命中: {query_param}")
            return cached_data
        
        # 生成唯一的任务 ID
        task_id = f"{self._get_query_key(query_param)}_{int(time.time() * 1000)}"
        
        # 创建结果事件和占位
        async with self.lock:
            self.results[task_id] = {"event": asyncio.Event(), "data": None}
        
        # 将任务放入队列
        await self.task_queue.put((task_id, query_param))
        self.logger.debug(f"任务已加入队列: {task_id}")
        # 通知处理协程有新任务
        self.event.set()
        
        # 等待结果，带超时
        try:
            self.logger.debug(f"等待任务结果: {task_id}, 超时: {timeout}秒")
            await asyncio.wait_for(self.results[task_id]["event"].wait(), timeout)
        except asyncio.TimeoutError:
            async with self.lock:
                # 清理
                if task_id in self.results:
                    del self.results[task_id]
            self.logger.error(f"获取数据超时: {query_param}")
            raise TimeoutError(f"获取数据超时: {query_param}")
        
        # 获取结果
        async with self.lock:
            result = self.results[task_id]["data"]
            # 清理
            del self.results[task_id]
        
        return result
    
    @abstractmethod
    def _get_query_key(self, query_param: T) -> str:
        """
        从查询参数获取唯一键
        
        Args:
            query_param: 查询参数
            
        Returns:
            用于标识任务的唯一键
        """
        pass
    
    @abstractmethod
    async def _check_cache(self, query_param: T,  cache_ttl: int) -> Optional[R]:
        """
        检查缓存中是否有有效的数据
        
        Args:
            query_param: 查询参数
            
        Returns:
            如果缓存中有有效数据，则返回数据；否则返回 None
        """
        pass
    
    @abstractmethod
    async def _fetch_data(self, query_params: List[T], ) -> List[R]:
        """
        从外部数据源获取数据
        
        Args:
            query_params: 查询参数列表
            
        Returns:
            查询结果数据列表
        """
        pass
    
    @abstractmethod
    async def _save_to_cache(self, query_param: T, data: R) -> None:
        """
        将数据保存到缓存
        
        Args:
            query_param: 查询参数
            data: 要缓存的数据
        """
        pass
    
    async def _process_tasks(self):
        """
        后台协程函数，批量处理任务队列中的任务
        """
        self.logger.info("任务处理协程已启动")
        while self.running:
            batch = []
            start_time = time.time()
            
            # 收集一批任务
            while len(batch) < self.batch_size and time.time() - start_time < self.max_wait_time:
                try:
                    # 非阻塞方式获取任务，最多等待剩余的最大等待时间
                    remaining_time = max(0, self.max_wait_time - (time.time() - start_time))
                    
                    # 使用 asyncio.wait_for 设置超时
                    try:
                        task = await asyncio.wait_for(self.task_queue.get(), remaining_time)
                        batch.append(task)
                    except asyncio.TimeoutError:
                        # 队列为空或超时，跳出循环
                        break
                except Exception as e:
                    self.logger.error(f"获取任务时出错: {e}")
                    break
            
            # 如果没有任务，等待新任务通知
            if not batch:
                self.event.clear()
                await self.event.wait()
                continue
            
            # 批量处理任务
            self.logger.info(f"开始处理批次任务，数量: {len(batch)}")
            await self._batch_process(batch)
    
    async def _batch_process(self, batch: List[Tuple[str, T]]):
        """
        批量处理一组任务
        
        Args:
            batch: 任务列表，每个任务是 (task_id, query_param) 的元组
        """
        # 从外部数据源获取数据
        self.logger.debug(f"从数据源获取数据，参数数量: {len(batch)}")
        results = await self._fetch_data([param for _, param in batch])
        
        # 将结果存储到缓存并通知等待的请求
        for (task_id, query_param), result in zip(batch, results):
            # 存储到缓存
            await self._save_to_cache(query_param, result)
            
            # 通知等待的请求
            async with self.lock:
                if task_id in self.results:
                    self.results[task_id]["data"] = result
                    self.results[task_id]["event"].set()
            
            # 标记任务完成
            self.task_queue.task_done()
    
    async def shutdown(self):
        """
        关闭服务，停止后台协程
        """
        self.running = False
        self.event.set()  # 唤醒协程
        await self.httpx_client.aclose()
        
        # 取消后台任务
        if self.worker_task and not self.worker_task.done():
            self.worker_task.cancel()
            try:
                await self.worker_task
            except asyncio.CancelledError:
                pass
            
            
class ConcurrentAsyncDataService(Generic[T, R], ABC):
    """
    基于并发池的异步数据服务抽象基类
    实现了并发控制和速率限制逻辑，子类只需实现特定的数据获取和缓存操作
    
    泛型参数:
        T: 查询参数类型 (如 str 或 int)
        R: 结果数据类型 (通常是 Dict[str, Any])
    """
    
    def __init__(self,lang: str = "en",  max_concurrency: int = 5, rate_limit: float = 5.0):
        """
        初始化并发异步数据服务
        
        Args:
            max_concurrency: 最大并发请求数
            rate_limit: 每秒最大请求数
        """
        self.lang = lang
        self.task_queue = asyncio.Queue()  # 异步任务队列
        self.results = {}  # 存储任务结果的字典
        self.max_concurrency = max_concurrency
        self.rate_limit = rate_limit
        self.request_interval = 1.0 / rate_limit if rate_limit > 0 else 0
        self.last_request_time = 0
        self.lock = asyncio.Lock()  # 用于异步安全访问结果字典
        self.rate_limit_lock = asyncio.Lock()  # 用于速率限制
        self.semaphore = asyncio.Semaphore(max_concurrency)  # 用于控制并发数
        self.worker_tasks = []  # 存储所有工作协程
        
        # 标记服务是否正在运行
        self.running = False
        
        # 获取日志记录器
        self.logger = get_logger(self.__class__.__name__)
        self.init_http_client()
    
    def init_http_client(self):
        self.httpx_client = httpx.AsyncClient(timeout=60.0, proxy=os.environ.get("HTTP_PROXY"))
    
    async def initialize(self):
        """
        异步初始化服务，启动工作协程池
        """
        if not self.running:
            self.running = True
            # 启动工作协程池
            for i in range(self.max_concurrency):
                task = asyncio.create_task(self._worker(i))
                self.worker_tasks.append(task)
            self.logger.info(f"服务已初始化并启动，最大并发数: {self.max_concurrency}, 速率限制: {self.rate_limit}/秒")
        return self
    
    async def submit_task(self, query_param: T,  cache_ttl: int, timeout: float= 360.0) -> R:
        """
        提交一个任务并等待结果
        
        Args:
            query_param: 查询参数 (如标题或ID)
            cache_ttl: 缓存有效期（秒）
            timeout: 等待结果的最大时间（秒）
            
        Returns:
            查询结果数据
            
        Raises:
            TimeoutError: 如果等待超时
        """
        # 首先检查缓存中是否已有数据
        self.logger.debug(f"提交任务: {query_param}")
        cached_data = await self._check_cache(query_param, cache_ttl)
        if cached_data:
            self.logger.debug(f"缓存命中: {query_param}")
            return cached_data
        
        # 生成唯一的任务 ID
        task_id = f"{self._get_query_key(query_param)}_{int(time.time() * 1000)}"
        
        # 创建结果事件和占位
        async with self.lock:
            self.results[task_id] = {"event": asyncio.Event(), "data": None}
        
        # 将任务放入队列
        await self.task_queue.put((task_id, query_param))
        self.logger.debug(f"任务已加入队列: {task_id}")
        
        # 等待结果，带超时
        try:
            self.logger.debug(f"等待任务结果: {task_id}, 超时: {timeout}秒")
            await asyncio.wait_for(self.results[task_id]["event"].wait(), timeout)
        except asyncio.TimeoutError:
            async with self.lock:
                # 清理
                if task_id in self.results:
                    del self.results[task_id]
            self.logger.error(f"获取数据超时: {query_param}")
            raise TimeoutError(f"获取数据超时: {query_param}")
        
        # 获取结果
        async with self.lock:
            result = self.results[task_id]["data"]
            # 清理
            del self.results[task_id]
        
        self.logger.debug(f"任务完成: {task_id}")
        return result
    
    async def _worker(self, worker_id: int):
        """
        工作协程，从队列获取任务并处理
        
        Args:
            worker_id: 工作协程ID，用于日志
        """
        self.logger.info(f"工作协程 {worker_id} 已启动")
        
        while self.running:
            try:
                # 从队列获取任务
                task_id, query_param = await self.task_queue.get()
                
                try:
                    # 使用信号量控制并发
                    async with self.semaphore:
                        # 应用速率限制
                        await self._apply_rate_limit()
                        
                        self.logger.debug(f"工作协程 {worker_id} 开始处理任务: {task_id}")
                        
                        # 获取数据
                        result = await self._fetch_single_data(query_param)
                        
                        # 存储到缓存
                        await self._save_to_cache(query_param, result)
                        
                        # 通知等待的请求
                        async with self.lock:
                            if task_id in self.results:
                                self.results[task_id]["data"] = result
                                self.results[task_id]["event"].set()
                                self.logger.debug(f"工作协程 {worker_id} 已设置任务结果: {task_id}")
                except Exception as e:
                    self.logger.error(f"工作协程 {worker_id} 处理任务出错: {task_id}, {e}")
                    # 通知等待的请求（出错）
                    async with self.lock:
                        if task_id in self.results:
                            # 创建错误结果
                            error_result = self._create_error_result(query_param)
                            self.results[task_id]["data"] = error_result
                            self.results[task_id]["event"].set()
                
                # 标记任务完成
                self.task_queue.task_done()
                
            except asyncio.CancelledError:
                self.logger.info(f"工作协程 {worker_id} 被取消")
                break
            except Exception as e:
                self.logger.error(f"工作协程 {worker_id} 发生异常: {e}")
    
    async def _apply_rate_limit(self):
        """
        应用速率限制，确保请求不超过指定的速率
        """
        if self.rate_limit <= 0:
            return
            
        async with self.rate_limit_lock:
            current_time = time.time()
            elapsed = current_time - self.last_request_time
            
            if elapsed < self.request_interval:
                # 需要等待以满足速率限制
                wait_time = self.request_interval - elapsed
                self.logger.debug(f"应用速率限制，等待 {wait_time:.4f} 秒")
                await asyncio.sleep(wait_time)
            
            # 更新最后请求时间
            self.last_request_time = time.time()
    
    @abstractmethod
    def _get_query_key(self, query_param: T) -> str:
        """
        从查询参数获取唯一键
        
        Args:
            query_param: 查询参数
            
        Returns:
            用于标识任务的唯一键
        """
        pass
    
    @abstractmethod
    async def _check_cache(self, query_param: T, cache_ttl: int) -> Optional[R]:
        """
        检查缓存中是否有有效的数据
        
        Args:
            query_param: 查询参数
            cache_ttl: 缓存有效期（秒）
            
        Returns:
            如果缓存中有有效数据，则返回数据；否则返回 None
        """
        pass
    
    @abstractmethod
    async def _fetch_single_data(self, query_param: T, ) -> R:
        """
        从外部数据源获取单个数据
        
        Args:
            query_param: 查询参数
            
        Returns:
            查询结果数据
        """
        pass
    
    @abstractmethod
    def _create_error_result(self, query_param: T) -> R:
        """
        创建错误结果
        
        Args:
            query_param: 查询参数
            
        Returns:
            错误结果数据
        """
        pass
    
    @abstractmethod
    async def _save_to_cache(self, query_param: T, data: R) -> None:
        """
        将数据保存到缓存
        
        Args:
            query_param: 查询参数
            data: 要缓存的数据
        """
        pass
    
    async def shutdown(self):
        """
        关闭服务，停止所有工作协程
        """
        self.logger.info("正在关闭服务...")
        self.running = False
        
        # 取消所有工作协程
        for task in self.worker_tasks:
            if not task.done():
                task.cancel()
        
        # 等待所有工作协程完成
        if self.worker_tasks:
            await asyncio.gather(*self.worker_tasks, return_exceptions=True)
        
        # 关闭 HTTP 客户端
        await self.httpx_client.aclose()
        
        self.logger.info("服务已关闭")